package util;

import javax.activation.UnsupportedDataTypeException;
import java.io.File;
import java.io.IOException;
import java.net.JarURLConnection;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.jar.JarFile;

public class ClassUtils {

    private ClassUtils() {
    }

    public static List<Class<?>> getClasses(String packageName) {
        String packagePath = packageName.replaceAll("\\.", "/");
        URL url = Thread.currentThread().getContextClassLoader().getResource(packagePath);
        switch (Objects.requireNonNull(url).getProtocol()) {
            case "file":
                return getClassFromFile(url, packagePath);
            case "jar":
                return getClassFromJar(url, packagePath);
            default:
                throw new RuntimeException(new UnsupportedDataTypeException());
        }
    }

    private static List<Class<?>> getClassFromFile(URL url, String packagePath) {
        URI uri;
        try {
            uri = url.toURI(); // uri for Chinese path
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }

        List<File> classFiles = new ArrayList<>();
        getClassFiles(new File(uri.getPath()), classFiles);
        List<Class<?>> classList = new ArrayList<>();
        for (File file : classFiles) {
            String filePath = file.getAbsolutePath().replaceAll("\\\\", "/");
            filePath = filePath.substring(filePath.lastIndexOf(packagePath));
            filePath = filePath.replace(".class", "").replaceAll("/", ".");
            try {
                classList.add(Class.forName(filePath));
            } catch (ClassNotFoundException e) {
                throw new RuntimeException(e);
            }
        }
        return classList;
    }

    private static List<Class<?>> getClassFromJar(URL url, String packagePath) {
        JarURLConnection urlConnection;
        JarFile jarFile;
        try {
            urlConnection = (JarURLConnection) url.openConnection();
            jarFile = urlConnection.getJarFile();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        List<Class<?>> classList = new ArrayList<>();
        jarFile.stream().filter(jarEntry -> !jarEntry.isDirectory()).filter(jarEntry -> {
            String entryName = jarEntry.getName();
            return entryName.startsWith(packagePath) && entryName.endsWith(".class");
        }).forEach(jarEntry -> {
            String entryName = jarEntry.getName().replace(".class", "").replaceAll("/", ".");
            try {
                classList.add(Class.forName(entryName));
            } catch (ClassNotFoundException e) {
                throw new RuntimeException(e);
            }
        });
        return classList;
    }

    private static void getClassFiles(File dir, List<File> result) {
        if (dir.isDirectory()) {
            Arrays.stream(dir.listFiles()).forEach(file -> getClassFiles(file, result));
        } else {
            if (dir.getName().endsWith(".class")) {
                result.add(dir);
            }
        }
    }
}
